package ai.koog.agents.core.agent

import ai.koog.agents.core.agent.AIAgent.Companion.State
import ai.koog.agents.core.agent.AIAgent.Companion.State.NotStarted
import ai.koog.agents.core.agent.context.AIAgentContext
import ai.koog.agents.core.agent.context.element.AgentRunInfoContextElement
import ai.koog.agents.core.agent.entity.AIAgentStorageKey
import ai.koog.agents.core.agent.entity.AIAgentStrategy
import ai.koog.agents.core.feature.AIAgentPipeline
import io.github.oshai.kotlinlogging.KLogger
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlin.uuid.ExperimentalUuidApi
import kotlin.uuid.Uuid

/**
 * Abstract base class representing a single-use AI agent with state.
 *
 * This AI agent is designed to execute a specific long-running strategy only once, and provides API to monitor and manage it's state.
 *
 * It maintains internal states including its running status, whether it was started, its result (if available), and
 * the root context associated with its execution. The class enforces safe state transitions and provides
 * thread-safe operations via a mutex.
 *
 * @param Input the type of the input accepted by the agent.
 * @param Output the type of the output produced by the agent.
 * @param TContext the type of the context used during the agent's execution, extending [AIAgentContext].
 * @property strategy the execution strategy defining how the agent processes input and produces output.
 * @property logger the logger used for logging execution details and errors.
 * @property id unique identifier for the agent instance, autogenerated if not provided during instantiation.
 */
@OptIn(ExperimentalUuidApi::class)
public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgentContext>(
    public val strategy: AIAgentStrategy<Input, Output, TContext>,
    private val logger: KLogger,
    agentId: String? = null
) : AIAgent<Input, Output> {
    /**
     * A mutex used to synchronize access to the state of the agent. Ensures that only one coroutine
     * can modify or read the shared state of the agent at a time, preventing data races and ensuring
     * thread-safe operations.
     */
    private val agentStateMutex: Mutex = Mutex()

    private var state: State<Output> = NotStarted()

    /**
     * A unique identifier represented as a string. This identifier is lazily initialized.
     * If a value is present in `agentId`, it will use that; otherwise, it will generate
     * a random UUID string.
     */
    final override val id: String by lazy { agentId ?: Uuid.random().toString() }

    final override suspend fun getState(): State<Output> = agentStateMutex.withLock { state.copy() }

    /**
     * Represents the pipeline used by the AI agent for processing tasks or data.
     *
     * This abstract property defines the structure or sequence of operations
     * within the AI agent's pipeline. It serves as the core mechanism for
     * executing workflows, handling inputs, and generating outputs in the
     * AI agent's functionality.
     */
    public abstract val pipeline: AIAgentPipeline

    private var agentFeatures: Map<AIAgentStorageKey<*>, Any>? = null

    /**
     * Retrieves a feature associated with the given key from the current context.
     *
     * @param key The key of the feature to retrieve.
     * @return The feature associated with the specified key, or null if no such feature exists.
     */
    @Suppress("UNCHECKED_CAST")
    override fun <Feature : Any> feature(key: AIAgentStorageKey<Feature>): Feature? {
        if (agentFeatures == null) {
            throw IllegalStateException(
                "Features are not ready yet. " +
                    "Probably your AI Agent hasn't stil initialized it's context. " +
                    "Have you run your agent?"
            )
        }

        return agentFeatures!![key] as? Feature
    }

    /**
     * Executes the agent's main functionality, coordinating with various components
     * such as pipelines and strategies. Ensures the agent is run in a thread-safe
     * manner using locking mechanisms.
     *
     * @param agentInput The input required to execute the agent's strategy.
     *                   This includes any data necessary for processing.
     * @return The output generated by the agent's execution, produced as a
     *         result of applying the strategy to the provided input.
     * @throws IllegalStateException if the agent was already started.
     * @throws Throwable if any exception occurs during the execution process.
     */
    final override suspend fun run(agentInput: Input): Output {
        agentStateMutex.withLock {
            if (state !is NotStarted) {
                throw IllegalStateException("Agent was already started")
            }
            state = State.Starting()
        }

        val runId = Uuid.random().toString()

        pipeline.prepareFeatures()

        return withContext(
            AgentRunInfoContextElement(
                agentId = this@StatefulSingleUseAIAgent.id,
                runId = runId,
                agentConfig = agentConfig,
                strategyName = strategy.name
            )
        ) {
            val context = prepareContext(agentInput, runId)

            agentFeatures = pipeline.getAgentFeatures(context)

            agentStateMutex.withLock {
                state = State.Running(context)
            }

            logger.debug {
                formatLog(
                    agentId = this@StatefulSingleUseAIAgent.id,
                    runId = runId,
                    message = "Starting agent execution"
                )
            }

            pipeline.onAgentStarting<Input, Output>(
                runId = runId,
                agent = this@StatefulSingleUseAIAgent,
                context = context
            )

            val result = try {
                strategy.execute(context = context, input = agentInput)
            } catch (e: Throwable) {
                logger.error(e) { "Execution exception reported by server!" }
                pipeline.onAgentExecutionFailed(
                    agentId = this@StatefulSingleUseAIAgent.id,
                    runId = runId,
                    throwable = e
                )
                agentStateMutex.withLock { state = State.Failed(e) }
                throw e
            }

            logger.debug {
                formatLog(
                    agentId = this@StatefulSingleUseAIAgent.id,
                    runId = runId,
                    message = "Finished agent execution"
                )
            }
            pipeline.onAgentCompleted(
                agentId = this@StatefulSingleUseAIAgent.id,
                runId = runId,
                result = result
            )

            agentStateMutex.withLock {
                state = if (result != null) {
                    State.Finished(result)
                } else {
                    State.Failed(Exception("result is null"))
                }
            }

            return@withContext result ?: error("result is null")
        }
    }

    /**
     * Closes the AI Agent and performs necessary cleanup operations.
     *
     * This method is a suspending function that ensures that the AI Agent's resources are released
     * when it is no longer needed. It notifies the pipeline of the agent's closure and ensures
     * that any associated features or stream providers are properly closed.
     *
     * Overrides the `close` method to implement agent-specific shutdown logic.
     */
    final override suspend fun close() {
        pipeline.onAgentClosing(agentId = this@StatefulSingleUseAIAgent.id)
        pipeline.closeFeaturesStreamProviders()
    }

    /**
     * Prepares and initializes the agent context required to handle the given input and run ID.
     *
     * @param agentInput the input provided to the agent for processing.
     * @param runId a unique identifier representing the current execution or operation run.
     * @return the initialized context specific to the agent setup for the provided input and run ID.
     */
    public abstract suspend fun prepareContext(agentInput: Input, runId: String): TContext

    /**
     * Formats a log message with the specified agent ID, run ID, and message content.
     *
     * @param agentId The unique identifier of the agent generating the log.
     * @param runId The unique identifier of the specific run or task associated with the log.
     * @param message The content of the log message to be formatted.
     * @return A formatted log string containing the agent ID, run ID, and the provided message.
     */
    protected fun formatLog(agentId: String, runId: String, message: String): String =
        "[agent id: $agentId, run id: $runId] $message"
}
